#include "lidar_localization/matching/localization_flow.hpp"

#include "glog/logging.h"
#include "lidar_localization/global_defination/global_defination.h"
#include <random>
#include "lidar_localization/tools/file_manager.hpp"
#include "lidar_localization/models/graph_optimizer/g2o/edge/edge_imu.hpp"
#include "lidar_localization/models/graph_optimizer/g2o/edge/edge_match.hpp"
#include "lidar_localization/models/graph_optimizer/g2o/edge/edge_odom.hpp"
namespace lidar_localization {
LocalizationFlow::LocalizationFlow(ros::NodeHandle& nh) {
    // subscriber
    cloud_sub_ptr_ = std::make_shared<CloudSubscriber>(nh, "/synced_cloud", 100000);
    imu_sub_ptr_ = std::make_shared<IMUSubscriber>(nh, "/kitti/oxts/imu", 1000000);
    laser_odom_sub_ptr_ = std::make_shared<OdometrySubscriber>(nh, "laser_odom", 1000000);
    gnss_pose_sub_ptr_ = std::make_shared<OdometrySubscriber>(nh, "/synced_gnss", 1000000);
    lidar_to_imu_ptr_ = std::make_shared<TFListener>(nh, "/imu_link", "/velo_link");

    // publisher
    global_map_pub_ptr_ = std::make_shared<CloudPublisher>(nh, "/global_map", "/map", 100);
    local_map_pub_ptr_ = std::make_shared<CloudPublisher>(nh, "/local_map", "/map", 100);
    current_scan_pub_ptr_ = std::make_shared<CloudPublisher>(nh, "/current_scan", "/map", 100);
    laser_localization_pub_ptr_ = std::make_shared<OdometryPublisher>(nh, "/laser_localization", "/map", "/lidar", 100);
    laser_tf_pub_ptr_ = std::make_shared<TFBroadCaster>("/map", "/vehicle_link");

    matching_ptr_ = std::make_shared<Matching>();
    std::string config_file_path = WORK_SPACE_PATH + "/config/matching/localization.yaml";
    YAML::Node config_node = YAML::LoadFile(config_file_path);
    odom_edge_noise_.resize(6);
    matching_noise_.resize(6);
    imu_noise_.resize(12);
    gravity_.resize(3);
    for (int i = 0; i < 6; ++i) {
        odom_edge_noise_(i) = config_node["odom_edge_noise"][i].as<double>();
        matching_noise_(i) = config_node["matching_noise"][i].as<double>();
    }
    for (int i = 0; i < 12; i++) {
        imu_noise_(i) = config_node["imu_noise"][i].as<double>();
    }
    for (int i = 0; i < 3; i++) {
        gravity_(i) = config_node["gravity"][i].as<double>();
    }
    key_distance_threshold_ = config_node["key_distance_threshold"].as<double>();
    key_time_threshold_ = config_node["key_time_threshold"].as<double>();
    window_size_ = config_node["window_size"].as<int>();
    std::string graph_optimizer_type = config_node["graph_optimizer_type"].as<std::string>();
    if (graph_optimizer_type == "g2o") {
        graph_optimizer_ptr_ = std::make_shared<G2oGraphOptimizer>("lm_var");
    } else {
        LOG(ERROR) << "没有找到与 " << graph_optimizer_type << " 对应的图优化模式,请检查配置文件";
    }
    std::cout << "后端优化选择的优化器为：" << graph_optimizer_type << std::endl << std::endl;

    FileManager::CreateDirectory(WORK_SPACE_PATH + "/localization_data");
    FileManager::CreateFile(ground_truth_ofs_, WORK_SPACE_PATH + "/localization_data/ground_truth.txt");
    FileManager::CreateFile(localization_ofs_, WORK_SPACE_PATH + "/localization_data/localization.txt");
}

bool LocalizationFlow::Run() {
    if (matching_ptr_->HasNewGlobalMap() && global_map_pub_ptr_->HasSubscribers()) {
        CloudData::CLOUD_PTR global_map_ptr(new CloudData::CLOUD());
        matching_ptr_->GetGlobalMap(global_map_ptr);
        global_map_pub_ptr_->Publish(global_map_ptr);
    }

    if (matching_ptr_->HasNewLocalMap() && local_map_pub_ptr_->HasSubscribers())
        local_map_pub_ptr_->Publish(matching_ptr_->GetLocalMap());

    if (!ReadData())
        return false;

    if (!InitCalibration()) 
        return false;

    if (!InitPose())
        return false;

    while(SyncData(true)) {
        Fuse();
        PublishData();
    }

    return true;
}

bool LocalizationFlow::ReadData() {
    static bool sensor_inited = false;
    cloud_sub_ptr_->ParseData(cloud_data_buff_);
    imu_sub_ptr_->ParseData(imu_data_buff_);
    laser_odom_sub_ptr_->ParseData(laser_odom_data_buff_);
    gnss_pose_sub_ptr_->ParseData(gnss_pose_data_buff_);

    if (cloud_data_buff_.empty() || imu_data_buff_.empty()
    || laser_odom_data_buff_.empty() || gnss_pose_data_buff_.empty())
        return false;

    if (!sensor_inited)
    {
        while (!cloud_data_buff_.empty())
        {
            if (imu_data_buff_.front().time > cloud_data_buff_.front().time
            || laser_odom_data_buff_.front().time > cloud_data_buff_.front().time
            || gnss_pose_data_buff_.front().time > cloud_data_buff_.front().time)
            {
                cloud_data_buff_.pop_front();
            }
            else
            {
                sensor_inited = true;
                break;
            }
        }
    }

    return sensor_inited;
}

bool LocalizationFlow::InitCalibration() {
    static bool calibration_received = false;
    if (!calibration_received) {
        if (lidar_to_imu_ptr_->LookupData(lidar_to_imu_)) {
            calibration_received = true;
        }
    }

    return calibration_received;
}

bool LocalizationFlow::InitPose() {
    static bool pose_inited = false;
    if (pose_inited)
    {
        return true;
    }
    if (!SyncData(false))
    {
        return false;
    }
    Eigen::Matrix4f laser_odometry;
    laser_odometry = current_gnss_pose_data_.pose * lidar_to_imu_;
    matching_ptr_->SetGNSSPose(laser_odometry);
    matching_ptr_->Update(current_cloud_data_, laser_odometry);
    matching_ptr_->TransformCurrentScan(current_cloud_data_, laser_odometry);
    laser_odometry = laser_odometry * lidar_to_imu_.inverse();
    current_pose_.time = current_cloud_data_.time;
    current_pose_.p = laser_odometry.block<3, 1>(0, 3).cast<double>();
    current_pose_.q = Eigen::Quaterniond(laser_odometry.block<3, 3>(0, 0).cast<double>());
    current_pose_.v = Eigen::Vector3d::Zero();
    current_pose_.ba = Eigen::Vector3d::Zero();
    current_pose_.bg = Eigen::Vector3d::Zero();
    pose_inited = true;
    graph_optimizer_ptr_->AddNode(current_pose_.p, current_pose_.q, current_pose_.v, current_pose_.ba, current_pose_.bg, false);
    graph_optimizer_ptr_->AddEdgeMatch(0, current_pose_.p, current_pose_.q, matching_noise_);
    key_frame_cnt_ = 0;
    Marginalize(true);
    PublishData();
    return true;
}

bool LocalizationFlow::SyncData(bool inited)
{
    if (cloud_data_buff_.empty())
    {
        return false;
    }
    current_cloud_data_ = cloud_data_buff_.front();
    double sync_time = current_cloud_data_.time;
    while (gnss_pose_data_buff_.size() > 1)
    {
        if (gnss_pose_data_buff_[1].time < sync_time)
        {
            gnss_pose_data_buff_.pop_front();
        }
        else
        {
            break;
        }
    }

    if (gnss_pose_data_buff_.size() > 1)
    {
        PoseData front_data = gnss_pose_data_buff_.at(0);
        PoseData back_data = gnss_pose_data_buff_.at(1);
        double front_scale = (back_data.time - sync_time) / (back_data.time - front_data.time);
        double back_scale = (sync_time - front_data.time) / (back_data.time - front_data.time);
        current_gnss_pose_data_.time = sync_time;
        current_gnss_pose_data_.pose.setIdentity();
        current_gnss_pose_data_.pose.block<3, 1>(0, 3) = front_data.pose.block<3, 1>(0, 3) * front_scale + back_data.pose.block<3, 1>(0, 3) * back_scale;
        current_gnss_pose_data_.pose.block<3,3>(0,0) = (front_data.GetQuaternion().slerp(back_scale, back_data.GetQuaternion())).toRotationMatrix();
    }
    else
    {
        return false;
    }

    while (laser_odom_data_buff_.size() > 1)
    {
        if (laser_odom_data_buff_[1].time < sync_time)
        {
            laser_odom_data_buff_.pop_front();
        }
        else
        {
            break;
        }
    }

    if (laser_odom_data_buff_.size() > 1)
    {
        PoseData front_data = laser_odom_data_buff_.at(0);
        PoseData back_data = laser_odom_data_buff_.at(1);
        double front_scale = (back_data.time - sync_time) / (back_data.time - front_data.time);
        double back_scale = (sync_time - front_data.time) / (back_data.time - front_data.time);
        current_laser_odom_data_.time = sync_time;
        current_laser_odom_data_.pose.setIdentity();
        current_laser_odom_data_.pose.block<3, 1>(0, 3) = front_data.pose.block<3, 1>(0, 3) * front_scale + back_data.pose.block<3, 1>(0, 3) * back_scale;
        current_laser_odom_data_.pose.block<3,3>(0,0) = (front_data.GetQuaternion().slerp(back_scale, back_data.GetQuaternion())).toRotationMatrix();
    }
    else
    {
        return false;
    }

    if (!inited)
    {
        current_imu_data_.clear();
        bool res = IMUData::SyncData(imu_data_buff_, current_imu_data_, sync_time);
        if (!res)
        {
            return false;
        }
    }
    else
    {
        bool res = IMUData::SyncData(imu_data_buff_, current_imu_data_, key_laser_odom_data_.time, sync_time);
        if (!res)
        {
            return false;
        }
    }
    cloud_data_buff_.pop_front();
    return true;
}

bool LocalizationFlow::PublishData() {
    Eigen::Matrix4f pose = Eigen::Matrix4f::Identity();
    pose.block<3, 3>(0, 0) = current_pose_.q.toRotationMatrix().cast<float>();
    pose.block<3, 1>(0, 3) = current_pose_.p.cast<float>();
    laser_tf_pub_ptr_->SendTransform(pose, current_cloud_data_.time);
    laser_localization_pub_ptr_->Publish(pose, current_cloud_data_.time);
    current_scan_pub_ptr_->Publish(matching_ptr_->GetCurrentScan());
    SavePose(ground_truth_ofs_, current_gnss_pose_data_.pose);
    SavePose(localization_ofs_, pose);
    return true;
}

bool LocalizationFlow::Fuse()
{
    Eigen::Matrix4d relative_pose = ((key_laser_odom_data_.pose * lidar_to_imu_.inverse()).inverse() * current_laser_odom_data_.pose * lidar_to_imu_.inverse()).cast<double>();
    Eigen::Vector3d pij = relative_pose.block<3, 1>(0, 3);
    Eigen::Quaterniond qij(relative_pose.block<3, 3>(0, 0));
    Eigen::Vector3d pj = state_.q * pij + state_.p;
    Eigen::Quaterniond qj = state_.q * qij;

    int node_num = graph_optimizer_ptr_->GetNodeNum();
    graph_optimizer_ptr_->AddNode(pj, qj, state_.v, state_.ba, state_.bg, false);
    
    graph_optimizer_ptr_->AddEdgeOdom(node_num - 5, node_num, pij, qij, odom_edge_noise_);

    Eigen::Matrix4f pose = Eigen::Matrix4f::Identity();
    pose.block<3, 3>(0, 0) = qj.toRotationMatrix().cast<float>();
    pose.block<3, 1>(0, 3) = pj.cast<float>();
    pose = pose * lidar_to_imu_;
    matching_ptr_->Update(current_cloud_data_, pose);
    pose = pose * lidar_to_imu_.inverse();
    pj = pose.block<3, 1>(0, 3).cast<double>();
    qj = Eigen::Quaterniond(pose.block<3, 3>(0, 0).cast<double>());
    graph_optimizer_ptr_->AddEdgeMatch(node_num, pj, qj, matching_noise_);

    std::shared_ptr<PreIntegration> pre_integration = std::make_shared<PreIntegration>(current_imu_data_, state_.ba, state_.bg, imu_noise_, gravity_);
    graph_optimizer_ptr_->AddEdgeImu(node_num - 5, node_num, pre_integration);

    graph_optimizer_ptr_->Optimize();

    current_pose_ = graph_optimizer_ptr_->GetState(node_num);
    current_pose_.time = current_cloud_data_.time;
    pose.block<3, 3>(0, 0) = current_pose_.q.toRotationMatrix().cast<float>();
    pose.block<3, 1>(0, 3) = current_pose_.p.cast<float>();
    pose = pose * lidar_to_imu_;
    matching_ptr_->TransformCurrentScan(current_cloud_data_, pose);

    Marginalize(false);

    return true;
}

void LocalizationFlow::Marginalize(bool first_key_frame)
{
    if (first_key_frame || current_pose_.time - state_.time > key_time_threshold_ || (current_pose_.p - state_.p).norm() > key_distance_threshold_)
    {
        key_frame_cnt_++;
        state_ = current_pose_;
        key_laser_odom_data_ = current_laser_odom_data_;
    }
    else
    {
        graph_optimizer_ptr_->DeleteNewestState();
    }
    
    if (key_frame_cnt_ > window_size_)
    {
        graph_optimizer_ptr_->Marginalize();
        key_frame_cnt_--;
    }

}

bool LocalizationFlow::SavePose(std::ofstream& ofs, const Eigen::Matrix4f& pose) {
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 4; ++j) {
            ofs << pose(i, j);
            
            if (i == 2 && j == 3) {
                ofs << std::endl;
            } else {
                ofs << " ";
            }
        }
    }

    return true;
}

}